"""
资源计算器

实现核心的资源分配算法，将人月数转换为整数人数分配
"""

import logging
from typing import List, Dict, Set
from datetime import date, timedelta
from dateutil.relativedelta import relativedelta
from collections import defaultdict

import sys
import os
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from models.simple_models import PhaseDetail, MonthlyAllocation


class ResourceCalculator:
    """资源计算器类"""

    def __init__(self):
        self.logger = logging.getLogger(__name__)

    def calculate_monthly_allocations(self, phases_detail: List[PhaseDetail],
                                    tl_enabled: bool = True,
                                    tl_ratio: int = 3) -> List[MonthlyAllocation]:
        """
        计算月度资源分配

        Args:
            phases_detail: 阶段详情列表
            tl_enabled: 是否启用TL计算
            tl_ratio: TL比例（每N个人配一个TL）

        Returns:
            List[MonthlyAllocation]: 月度分配列表
        """
        self.logger.info("开始计算月度资源分配...")

        allocations = []

        # 按阶段处理
        for phase in phases_detail:
            phase_allocations = self._calculate_phase_allocations(phase)
            allocations.extend(phase_allocations)

        # 如果启用TL计算，添加TL分配
        if tl_enabled:
            tl_allocations = self._calculate_tl_allocations(allocations, tl_ratio)
            allocations.extend(tl_allocations)

        # 按月份和角色合并重复分配
        merged_allocations = self._merge_allocations(allocations)

        self.logger.info(f"计算完成，共生成 {len(merged_allocations)} 个月度分配")

        return merged_allocations

    def _calculate_phase_allocations(self, phase: PhaseDetail) -> List[MonthlyAllocation]:
        """
        计算单个阶段的资源分配

        Args:
            phase: 阶段详情

        Returns:
            List[MonthlyAllocation]: 分配列表
        """
        allocations = []

        # 计算阶段持续的月份数
        duration_months = self._calculate_duration_months(phase.start_date, phase.end_date)

        if duration_months == 0:
            self.logger.warning(f"阶段 {phase.phase_name} 持续时间为0，跳过")
            return allocations

        # 使用整数分配算法
        person_counts = self._allocate_person_months(phase.person_months, duration_months)

        # 生成每个月的分配记录
        current_date = phase.start_date.replace(day=1)  # 月初
        for i, person_count in enumerate(person_counts):
            if person_count > 0:  # 只记录有人员分配的月份
                allocation = MonthlyAllocation(
                    project_id=phase.project_id,
                    phase_name=phase.phase_name,
                    month=current_date,
                    role=phase.role,
                    person_count=person_count,
                    unit_price=phase.unit_price,
                    cost=person_count * phase.unit_price
                )
                allocations.append(allocation)

            # 移动到下个月
            current_date += relativedelta(months=1)

        return allocations

    def _calculate_duration_months(self, start_date: date, end_date: date) -> int:
        """
        计算持续月份数

        Args:
            start_date: 开始日期
            end_date: 结束日期

        Returns:
            int: 持续月份数
        """
        # 计算跨越的月份数
        start_month = start_date.replace(day=1)
        end_month = end_date.replace(day=1)

        months = 0
        current = start_month
        while current <= end_month:
            months += 1
            current += relativedelta(months=1)

        return months

    def _allocate_person_months(self, total_person_months: float, duration_months: int) -> List[int]:
        """
        整数分配算法：将人月数分配到各个月份

        Args:
            total_person_months: 总人月数
            duration_months: 持续月份数

        Returns:
            List[int]: 每个月的人数分配
        """
        if duration_months <= 0:
            return []

        # 基础分配：每个月分配的基本人数
        base_allocation = int(total_person_months // duration_months)

        # 剩余的人月数
        remainder = total_person_months - (base_allocation * duration_months)

        # 初始化分配数组
        allocations = [base_allocation] * duration_months

        # 将剩余的人月数按优先级分配
        # 优先分配给前面的月份
        remainder_int = int(remainder)
        for i in range(min(remainder_int, duration_months)):
            allocations[i] += 1

        self.logger.debug(f"分配 {total_person_months} 人月到 {duration_months} 个月: {allocations}")

        return allocations

    def _calculate_tl_allocations(self, base_allocations: List[MonthlyAllocation],
                                 tl_ratio: int) -> List[MonthlyAllocation]:
        """
        计算TL（Team Leader）分配

        Args:
            base_allocations: 基础分配列表
            tl_ratio: TL比例（每N个人配一个TL）

        Returns:
            List[MonthlyAllocation]: TL分配列表
        """
        tl_allocations = []

        # 按项目和月份统计非TL人员数量
        monthly_stats = defaultdict(lambda: {'total_persons': 0, 'projects': set()})

        for alloc in base_allocations:
            if alloc.role != 'TL':  # 排除已有的TL
                key = (alloc.month, alloc.project_id)
                monthly_stats[key]['total_persons'] += alloc.person_count
                monthly_stats[key]['projects'].add(alloc.project_id)

        # 计算每个月需要的TL数量
        for (month, project_id), stats in monthly_stats.items():
            total_persons = stats['total_persons']

            if total_persons > 0:
                # 计算需要的TL数量（每tl_ratio个人配一个TL，至少1个）
                tl_count = max(1, int((total_persons + tl_ratio - 1) // tl_ratio))

                # 查找该项目的TL单价（如果没有，使用默认值）
                tl_unit_price = self._get_tl_unit_price(base_allocations, project_id)

                # 创建TL分配记录
                tl_allocation = MonthlyAllocation(
                    project_id=project_id,
                    phase_name="管理工作",  # TL的工作描述
                    month=month,
                    role="TL",
                    person_count=tl_count,
                    unit_price=tl_unit_price,
                    cost=tl_count * tl_unit_price
                )
                tl_allocations.append(tl_allocation)

                self.logger.debug(f"为 {project_id} 的 {month.strftime('%Y-%m')} "
                                f"分配 {tl_count} 个TL（基于 {total_persons} 个团队成员）")

        self.logger.info(f"计算TL分配完成，共生成 {len(tl_allocations)} 个TL分配")

        return tl_allocations

    def _get_tl_unit_price(self, base_allocations: List[MonthlyAllocation], project_id: str) -> float:
        """
        获取TL的单价

        Args:
            base_allocations: 基础分配列表
            project_id: 项目ID

        Returns:
            float: TL单价
        """
        # 尝试从现有分配中查找TL单价
        for alloc in base_allocations:
            if alloc.project_id == project_id and alloc.role == 'TL':
                return alloc.unit_price

        # 如果没有找到，使用SE单价的1.2倍作为默认值
        se_price = None
        for alloc in base_allocations:
            if alloc.project_id == project_id and alloc.role == 'SE':
                se_price = alloc.unit_price
                break

        if se_price:
            return se_price * 1.2

        # 如果还是没有，使用默认值
        return 700.0

    def _merge_allocations(self, allocations: List[MonthlyAllocation]) -> List[MonthlyAllocation]:
        """
        合并重复的分配记录

        Args:
            allocations: 原始分配列表

        Returns:
            List[MonthlyAllocation]: 合并后的分配列表
        """
        # 按 (项目ID, 月份, 角色) 分组
        grouped = defaultdict(list)

        for alloc in allocations:
            key = (alloc.project_id, alloc.month, alloc.role)
            grouped[key].append(alloc)

        merged = []

        for (project_id, month, role), group in grouped.items():
            if len(group) == 1:
                # 只有一个记录，直接添加
                merged.append(group[0])
            else:
                # 多个记录，需要合并
                total_person_count = sum(alloc.person_count for alloc in group)
                total_cost = sum(alloc.cost for alloc in group)

                # 使用第一个记录的基本信息，但更新数量和成本
                base_alloc = group[0]

                # 合并阶段名称
                phase_names = list(set(alloc.phase_name for alloc in group))
                merged_phase_name = "/".join(sorted(phase_names))

                merged_alloc = MonthlyAllocation(
                    project_id=project_id,
                    phase_name=merged_phase_name,
                    month=month,
                    role=role,
                    person_count=total_person_count,
                    unit_price=base_alloc.unit_price,
                    cost=total_cost
                )
                merged.append(merged_alloc)

                self.logger.debug(f"合并 {len(group)} 个分配记录: "
                                f"{project_id}-{month.strftime('%Y-%m')}-{role}")

        # 按月份和角色排序
        merged.sort(key=lambda x: (x.month, x.role, x.project_id))

        return merged